from langchain_openai import ChatOpenAI
from langchain_openai import OpenAIEmbeddings
from langchain_community.embeddings import ZhipuAIEmbeddings

import os

model = ChatOpenAI(model="deepseek-chat",
                   api_key=os.environ.get('DEEPSEEK_API_KEY'),
                   base_url=os.environ.get('DEEPSEEK_BASE_URL'))

# embed = OpenAIEmbeddings(
#     api_key='sk-proj-Vs1hvbqAjlKBnTS9yTrNygLSq16ElNdDuTV1ADDgfezt3b4B5oA',
#     base_url='https://ai.devtool.tech/proxy/v1',     # 国内中转地址
#     model="text-embedding-3-large"
# )

embed = ZhipuAIEmbeddings(
    #api_key=os.environ.get('ZHIPU_API_KEY_1'),
    api_key='0d28f030249b4fe38dc501510748b595.9SGm9tuJBlcgqKBm',
    model="embedding-3"
)

from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document


# 加载一些模拟的假数据
doc1 = Document(page_content="员工每年享有一定数量的病假。有关资格和具体细节可以在员工手册中找到。")
doc2 = Document(page_content="员工请病假时，必须首先通知其主管关于病情和预计缺勤时间。员工需填写病假申请表，并提交给人力资源部门或主管。")
doc3 = Document(page_content="病假申请表可以在公司内部网找到。表格需要填写员工姓名、部门、缺勤日期和缺勤原因等信息。")


# 创建 Faiss 向量存储
vector_store = FAISS.from_documents([doc1, doc2, doc3], embed)

# 将文件保存到本地，包括：向量数据、索引文件和元数据文件
vector_store.save_local(folder_path='../FAISS_Data')

# 加载本地的Faiss向量文件，allow_dangerous_deserialization 用于控制是否允许在加载向量存储时进行潜在的危险反序列化操作。
vector_store = FAISS.load_local(embeddings=embed, folder_path='../FAISS_Data', allow_dangerous_deserialization=True)

# 将 FAISS 向量存储转换为一个 retriever（检索器），并为该检索器设置一些搜索相关的参数。k=1 表示检索时返回 最相似的 1 个文档
retriever = vector_store.as_retriever(search_kwargs={'k': 1})

# 执行相似度搜素
query = "请问我们公司有没有病假？"
results = retriever.invoke(query)

for doc in results:
    print(f"Content: {doc.page_content}")